import { create } from "zustand";
import { immer } from "zustand/middleware/immer";
import { convertFileToBase64 } from "../../util/file";
import { persist } from "zustand/middleware";
import { zustandPersistStorage } from "../../util/storage";
import { isImageFile } from "../shared/utils/file_utils";
import { useChat } from "./useChat";

export type FileStatus = "uploading" | "success" | "error";

export interface UploadFile {
    id: string;
    // which conversation message use this file
    conversationId?: string;
    messageId?: string;
    progress: number;
    status: FileStatus;
    base64: string;
    isImage: boolean;
    fileName?: string;
    fileSize?: string;
}

interface UploadState {
    files: UploadFile[];
}

interface UploadActions {
    addFiles: (files: File[]) => void;
    removeFile: (id: string) => void;
    removeConversationFiles: (conversationId: string) => void;
    markFileUsed: (fileId: string, messageId: string) => void;
    updateProgress: (id: string, progress: number) => void;
    updateStatus: (id: string, status: FileStatus) => void;
    clearFiles: () => void;
    getSupportedTypes: (capabilities: string[]) => string[] | undefined;
}

type UploadStore = UploadState & UploadActions;

export const IMAGE_TYPES = ["image/jpeg", "image/png", "image/gif"];
export const DOCUMENT_TYPES = [
    "application/pdf",
    "text/plain",
    "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
    "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
    "text/csv",
    "application/vnd.openxmlformats-officedocument.presentationml.presentation",
];

/**
 * 文件上传状态管理器
 * - 管理上传文件列表
 * - 管理文件上传进度
 * - 管理文件预览
 */
const uploadStore = immer<UploadStore>((set, get) => ({
    files: [],
    getSupportedTypes: (capabilities) => {
        if (capabilities.includes("vision")) {
            return IMAGE_TYPES;
        }
        if (capabilities.includes("multimodal")) {
            return [...IMAGE_TYPES, ...DOCUMENT_TYPES];
        }
    },
    addFiles: (newFiles) => {
        const chatStore = useChat.getState();
        const currentConversationId = chatStore.currentConversationId;
        // 为每个文件创建预览URL
        newFiles.forEach(async (file) => {
            const base64 = await convertFileToBase64(file);
            set((state) => {
                const id = crypto.randomUUID();
                const isImage = isImageFile(file.type);
                const fileObj: UploadFile = {
                    id,
                    conversationId: currentConversationId,
                    progress: 100,
                    status: "success",
                    base64,
                    isImage,
                    fileName: file.name,
                    fileSize: file.size
                        ? (file.size / 1024).toFixed(2) + " KB"
                        : undefined,
                };
                state.files.push(fileObj);
            });
        });
    },
    markFileUsed: (fileId, messageId) => {
        set((state) => {
            const file = state.files.find((f) => f.id === fileId);
            if (!file) {
                throw new Error("Mark a non-exist file as used: " + fileId);
            }
            file.messageId = messageId;
        });
    },
    removeConversationFiles: (conversationId) => {
        set((state) => {
            state.files = state.files.filter(
                (f) => f.conversationId !== conversationId,
            );
        });
    },
    removeFile: (id) => {
        set((state) => {
            state.files = state.files.filter((f) => f.id !== id);
        });
    },
    updateProgress: (id, progress) => {
        set((state) => {
            const file = state.files.find((f) => f.id === id);
            if (file) {
                file.progress = progress;
            }
        });
    },
    updateStatus: (id, status) =>
        set((state) => {
            const file = state.files.find((f) => f.id === id);
            if (file) {
                file.status = status;
            }
        }),

    clearFiles: () =>
        set((state) => {
            state.files = [];
        }),
}));

export const useUpload = create<UploadStore>()(
    persist(uploadStore, {
        name: "upload_store",
        storage: zustandPersistStorage,
        partialize: (state) => ({ files: state.files }),
    }),
);
